import os, cv2
import numpy as np
from PyQt5.QtGui import QPixmap
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import keras
import tensorflow as tf

model = keras.models.load_model('mnist.h5')

def QPixmap2Mat(img : QPixmap):
    qimg = img.toImage()
    temp_shape = (qimg.height(), qimg.bytesPerLine() * 8 // qimg.depth())
    temp_shape += (4,)
    ptr = qimg.bits()
    ptr.setsize(qimg.byteCount())
    result = np.array(ptr, dtype=np.uint8).reshape(temp_shape)
    result = result[..., :3]
    return result

def expandImage(img):
    w, h = img.shape[0 : 2]
    size = max(w, h)
    x = (size - w) // 2
    y = (size - h) // 2
    image = np.zeros((size, size), dtype='uint8')
    for i in range(w):
        for j in range(h):
            image[i + x, j + y] = img[i, j]
    return image

def predictImage(img):
    img = cv2.resize(img, (28, 28))
    img = np.expand_dims(img, axis=0)
    return np.argmax(model.predict(img))

def predict(img: QPixmap):
    img = cv2.cvtColor(QPixmap2Mat(img), cv2.COLOR_RGB2GRAY)
    contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    string = ''
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        number = img[y : y + h, x : x + w]
        string += str(predictImage(expandImage(number)))
    return ' '.join(string)

